import pandas as pd
from sqlalchemy import create_engine, text, MetaData

class MirrorDB:

    def __init__(self, db_url):
        """
        Initialize the class with database URL and table name.

        :param db_url: MySQL database URL (e.g., 'mysql+pymysql://user:password@host:port/database')
        :param table_name: Name of the table to store the CSV data.
        """
        self.db_url = db_url
        self.engine = create_engine(db_url)

        # 反射元数据
        metadata = MetaData()
        metadata.reflect(bind=self.engine)

        # 获取所有表名
        self.tables = list(metadata.tables.keys())

        print("数据库中的所有表：")
        for table in self.tables:
            print(table)

    def csv_to_db(self, table_name, csv_file_path, tbindex=False): #如果使用csv文件就把参数里的df改为csv_file_path，同时取消第一行df注释
        """
        Write a CSV file to the MySQL database.

        :param csv_file_path: Path to the CSV file.
        """
        # Read the CSV file into a pandas DataFrame
        df = pd.read_csv(csv_file_path)
        return self.df_to_db(table_name, df, tbindex=False)


    def df_to_db(self,table_name, df, tbindex=False):
        """
        Write a CSV file to the MySQL database.

        :param csv_file_path: Path to the CSV file.
        """

       # Find columns that start with "Unnamed:"
        unnamed_columns = [col for col in df.columns if col.startswith("Unnamed:")]

        # Drop these columns if they exist
        if unnamed_columns:
            print(f"Dropping columns: {unnamed_columns}")
            df = df.drop(columns=unnamed_columns)

        # Write the DataFrame to the MySQL database
        df.to_sql(table_name, con=self.engine, if_exists='append', index=tbindex)
        print(f"Data from dataframe has been written to the table {table_name}.")

    def db_to_csv(self, querysql, csv_file_path):
        """
        Read data from the MySQL database and generate a CSV file.

        :param csv_file_path: Path to save the output CSV file.
        """
        # Read data from the MySQL database into a pandas DataFrame
        #query = f"SELECT * FROM {self.table_name}"
        df = pd.read_sql(querysql, con=self.engine)

        # Write the DataFrame to a CSV file
        df.to_csv(csv_file_path, index=False)
        print(f"Data from table {self.table_name} has been written to {csv_file_path}.")

    def close(self):
        """
        Close the database connection.
        """
        self.engine.dispose()
        print("Database connection closed.")

    def create_table_from_df(self,table_name, df, dropflag=False):
        """
        Create a MySQL table based on the structure of a dateframe.

        :param db_url: MySQL database URL (e.g., 'mysql+pymysql://user:password@host:port/database')
        :param table_name: Name of the table to create.
        :param df: dataframe.
        """

        #if table exists and drop flag is false, return
        if dropflag==False and table_name in self.tables:
            return

        # Read columns from a pandas DataFrame
        header_row = df.iloc[0]
        columns_to_drop = header_row.index[header_row.isna()]
        df.drop(columns=columns_to_drop, inplace=True)

        # Generate SQL data types based on the DataFrame's column types
        dtype_mapping = {
            'int64': 'INT',
            'float64': 'FLOAT',
            'object': 'VARCHAR(255)',  # Default to VARCHAR(255) for string columns
            'bool': 'BOOLEAN',
            'datetime64[ns]': 'DATETIME',
        }

        # Create a list of column definitions for the SQL CREATE TABLE statement
        columns = []
        for col_name, col_type in df.dtypes.items():
            if col_name.startswith("Unnamed:"):
                continue
            sql_type = dtype_mapping.get(str(col_type), 'VARCHAR(255)')  # Default to VARCHAR if type is unknown
            columns.append(f"`{col_name}` {sql_type}")

        # Generate the CREATE TABLE SQL statement
        create_table_sql = f"CREATE TABLE `{table_name}` (\n"
        create_table_sql += ",\n".join(columns)
        create_table_sql += "\n);"

        # Connect to the MySQL database and execute the CREATE TABLE statement

        with self.engine.connect() as connection:
            try:
                # Drop the table if it already exists
                connection.execute(text(f"DROP TABLE IF EXISTS `{table_name}`;"))
                # Create the new table
                connection.execute(text(create_table_sql))
                print(f"Table '{table_name}' created successfully.")
                self.tables.append(table_name)
            except Exception as e:
                print(f"Error creating table: {e}")
            finally:
                # Close the database connection
                self.engine.dispose()



def create_table_from_csv(db_url, table_name, csv_file_path):
    """
    Create a MySQL table based on the structure of a CSV file.

    :param db_url: MySQL database URL (e.g., 'mysql+pymysql://user:password@host:port/database')
    :param table_name: Name of the table to create.
    :param csv_file_path: Path to the CSV file.
    """
    # Read the CSV file into a pandas DataFrame
    df = pd.read_csv(csv_file_path)
    header_row = df.iloc[0]
    columns_to_drop = header_row.index[header_row.isna()]
    df.drop(columns=columns_to_drop, inplace=True)


    # Generate SQL data types based on the DataFrame's column types
    dtype_mapping = {
        'int64': 'INT',
        'float64': 'FLOAT',
        'object': 'VARCHAR(255)',  # Default to VARCHAR(255) for string columns
        'bool': 'BOOLEAN',
        'datetime64[ns]': 'DATETIME',
    }

    # Create a list of column definitions for the SQL CREATE TABLE statement
    columns = []
    for col_name, col_type in df.dtypes.items():
        sql_type = dtype_mapping.get(str(col_type), 'VARCHAR(255)')  # Default to VARCHAR if type is unknown
        columns.append(f"`{col_name}` {sql_type}")

    # Generate the CREATE TABLE SQL statement
    create_table_sql = f"CREATE TABLE `{table_name}` (\n"
    create_table_sql += ",\n".join(columns)
    create_table_sql += "\n);"

    # Connect to the MySQL database and execute the CREATE TABLE statement
    engine = create_engine(db_url)
    with engine.connect() as connection:
        try:
            # Drop the table if it already exists
            connection.execute(text(f"DROP TABLE IF EXISTS `{table_name}`;"))
            # Create the new table
            connection.execute(text(create_table_sql))
            print(f"Table '{table_name}' created successfully.")
        except Exception as e:
            print(f"Error creating table: {e}")
        finally:
            # Close the database connection
            engine.dispose()


# Example Usage
if __name__ == "__main__":
    # Database URL (replace with your MySQL credentials)
    #db_url = "mysql+pymysql://user:password@localhost:3306/test_db"
    db_url = "mysql+pymysql://user:password@ip:port/database"

    db1=MirrorDB(db_url)

    # Path to the input CSV file
    input_csv = "test.csv"

    df=pd.read_csv(input_csv)

    db1.create_table_from_df("daily_000001_SZ",df,True)

    db1.csv_to_db("daily_000001_SZ",input_csv)

